
import math

import pandas as pd
import altair as alt

import torch
import torch.nn as nn


class PositionEncodingLayer(nn.Module):
    # max len 是 整个 seq 最多有多少个 token demodel 是 一个位置编码向量是多长
    def __init__(self, d_model, max_len = 5000, dropout = 0.1):
        super(PositionEncodingLayer, self).__init__()

        self.dropout = nn.Dropout(dropout)

        # 一次性把所有的位置编码全都算出来，方便后面节省时间
        pe = torch.zeros(max_len, d_model)

        # 生成位置 [0, 1, 2, 3, 4, .., d_model - 1]
        # 然后将其从 (max_len, ) 升维到二维 (max_len, 1)
        pos = torch.arange(0, max_len).unsqueeze(1)

        # 指数运算在底数或者指数非常大或者非常小的时候不是很稳定，采用 log + exp 的方式更好
        # 计算 1 / (10000 ^ (2i / d_model)) , 也就是 sin(omega * pos) 中的 omega
        # torch.arange(0, d_model, 2) -> [0, 2, 4, 6, 8, 10, 12 , ...] 也就是 2i
        # 最终生成 (d_model // 2,) 形状的行向量
        omega = torch.exp((torch.arange(0, d_model, 2)) * (-math.log(10000) / d_model))

        # 生成偶数位置的数值 PE = sin (pos * omega)
        pe[:, 0::2] = torch.sin(pos * omega)
        # 生成奇数位置上的数值 PE = cos(pos * omega)
        pe[:, 1::2] = torch.cos(pos * omega)

        # 为 batch 维度预留位置，方便后面可以进行广播
        pe = pe.unsqueeze(0)

        # 注册成一个常量，不参与 参数更新
        self.register_buffer("pe", pe)

    # params [x]: size (batch, seq_len, d_model) input
    def forward(self, x):
        seq_len = x.size(1)
        # batch 维取全局， max_len 维从0取到seq_len, d_model 维取全部
        return self.dropout(x + self.pe[:, : seq_len, :].requires_grad_(False))


if __name__ == "__main__":
    # example_positional().show()
    # show_example(example_positional())

    def example_positional():
        pe = PositionEncodingLayer(20, 5000, 0)
        a = torch.zeros(1, 100, 20)
        y = pe(a)

        print(y)

        data = pd.concat(
            [
                pd.DataFrame(
                    {
                        "embedding": y[0, :, dim],
                        "dimension": dim,
                        "position": list(range(100)),
                    }
                )
                for dim in [4, 5, 6, 7]
            ]
        )

        return (
            alt.Chart(data)
            .mark_line()
            .properties(width=800)
            .encode(x="position", y="embedding", color="dimension:N")
            .interactive()
        )

    chart = example_positional()
    chart.save("positional_encoding.html")
